/*
 *  dag_teach.cpp
 *  dag
 *
 *  Created by Phillip Popp on 3/27/11.
 *  Copyright 2011 __MyCompanyName__. All rights reserved.
 *
 */

#include "dag_teach.h"
#include "dag_apply_link.h"

#include <map>
#include <set>
#include <vector>

using namespace dag;
using namespace std;

struct DAG_LOCAL DepotAndLink {
	FrameDepot* depot;
	const FrameLink* link;
};
static void teach_node(Node& node,
					   const Graph& graph,
					   set<Node*>& taughtNodes,
					   map<Node*, FrameDepot*>& inputDepots,
					   map<Node*, FrameDepot*>& outputDepots);
static void get_immediate_inputs(Node& node,
								 const Graph& graph,
								 set<Node*>& inputs);
static void get_immediate_outputs(Node& node,
								  const Graph& graph,
								  set<Node*>& inputs);
static void delete_frame_depots(map<Node*, FrameDepot*>& frames);
template <bool forward>
static void map_frame_depots(const vector<Edge*>& from,
							 FrameDepot& to,
							 map<Node*, FrameDepot*>& inputDepots);
static void prepare_inputs(Node& node,
						   const Graph& graph,
						   set<Node*>& taughtNodes,
						   map<Node*, FrameDepot*>& inputDepots,
						   map<Node*, FrameDepot*>& outputDepots);
static void prepare_outputs(Node& node,
							const Graph& graph,
							set<Node*>& taughtNodes,
							map<Node*, FrameDepot*>& inputDepots,
							map<Node*, FrameDepot*>& outputDepots);

static auto_ptr<FrameDepot> create_node_input(Node& node,
										   const Graph& graph,
										   set<Node*>& taughtNodes,
										   map<Node*, FrameDepot*>& inputDepots,
										   map<Node*, FrameDepot*>& outputDepots);
static auto_ptr<FrameDepot> create_node_output(Node& node,
											const Graph& graph,
											set<Node*>& taughtNodes,
											map<Node*, FrameDepot*>& inputDepots,
											map<Node*, FrameDepot*>& outputDepots);

void dag::teach_graph(const Graph& graph, 
				 const FrameDepot* input,
				 const FrameDepot* output) {

	/* validate inputs */
	NodeTraits traits = get_node_set_traits(graph.getNodes());
	switch (traits.resources) {
		case LearnResource_Inputs:
			if (input == NULL) {
				throw NullPtr();
			}
			break;
		case LearnResource_Outputs:
			if (output == NULL) {
				throw NullPtr();
			}
			break;
		case LearnResource_InputsOutputs:
			if (output == NULL || input == NULL) {
				throw NullPtr();
			}
			break;
	}
	
	
	/* intialize node-framedepot map with inputs and outputs */
	/* TODO: shouldn't need this stupid const_cast (on the nodes).  WTF is up
	 * with this compiler? */
	map<Node*, FrameDepot*> inputDepots;
	inputDepots[const_cast<Node*>(graph.getHeadNode())] = const_cast<FrameDepot*>(input);

	map<Node*, FrameDepot*> outputDepots;
	outputDepots[const_cast<Node*>(graph.getTailNode())] = const_cast<FrameDepot*>(output);

	/* teach nodes one by one */
	set<Node*> taughtNodes;
	set<Node*>::const_iterator nodeIter;
	const set<Node*> nodes = graph.getNodes();
	
	for (nodeIter = nodes.begin(); nodeIter != nodes.end(); nodeIter++) {
		if (*nodeIter == NULL) {
			throw NullPtr();
		}
		
		if (taughtNodes.count(*nodeIter) == 0) {
			teach_node(**nodeIter, 
					   graph, 
					   taughtNodes, 
					   inputDepots, 
					   outputDepots);
		}
	}
	
	/* delete temporary frame depots created during treaching, but don't delete 
	 * the given input and output depots */
	/* TODO: shouldn't need this stupid const_cast.  WTF is up with this 
	 * compiler? */
	inputDepots.erase(const_cast<Node*>(graph.getHeadNode()));
	outputDepots.erase(const_cast<Node*>(graph.getTailNode()));
	delete_frame_depots(inputDepots);
	delete_frame_depots(outputDepots);
}

/* teach a single node in the graph.  recursively teach other nodes if needed */
void teach_node(Node& node,
				const Graph& graph,
				set<Node*>& taughtNodes,
				map<Node*, FrameDepot*>& inputDepots,
				map<Node*, FrameDepot*>& outputDepots) {
	/* check if node is already taught */
	if (taughtNodes.count(&node) == 1) {
		return;
	}
	
	/* teach based off of learn resources */
	if (node.traits().resources == LearnResource_None) {
		node.learn(NULL, NULL);
	}
	else if (node.traits().resources == LearnResource_Inputs) {
		prepare_inputs(node,
					   graph,
					   taughtNodes,
					   inputDepots,
					   outputDepots);
		node.learn(inputDepots[&node], NULL);
	}
	else if (node.traits().resources == LearnResource_Outputs) {
		prepare_outputs(node,
						graph,
						taughtNodes,
						inputDepots,
						outputDepots);
		node.learn(NULL, outputDepots[&node]);
	}
	else if (node.traits().resources == LearnResource_InputsOutputs) {
		prepare_inputs(node,
					   graph,
					   taughtNodes,
					   inputDepots,
					   outputDepots);
		prepare_outputs(node,
						graph,
						taughtNodes,
						inputDepots,
						outputDepots);
		node.learn(inputDepots[&node], outputDepots[&node]);
	}
	else {
		throw DagError("unimplemented learn resource condition");
	}
	
	taughtNodes.insert(&node);
}

template <bool _forward>
void map_frame_depots(const vector<Edge*>& from,
					  FrameDepot& to,
					  map<Node*, FrameDepot*>& inputDepots) {
	
	/* retrieve and validate incoming frame depots */
	vector<DepotAndLink> fromDepots;
	size_t numFrames = 0;
	
	for (size_t i = 0; i < from.size(); i++) {
		if (from[i] == NULL) {
			/* TODO: better error reporting */
			throw NullPtr();
		}
		
		FrameDepot* depot = NULL;
		if (_forward) {
			depot = inputDepots[from[i]->from];
		}
		else {
			depot = inputDepots[from[i]->to];
		}
		
		if (depot == NULL) {
			/* TODO: better error reporting */
			throw DagError("missing frame depot");
		}
		
		if (i == 0) {
			numFrames = depot->size();
		}
		else {
			if (numFrames != depot->size()) {
				/* TODO: better error reporting */
				throw DagError("inconsistent number of frames in frame depot");
			}
		}
		
		DepotAndLink das;
		das.depot = depot;
		das.link = &from[i]->link;
		
		fromDepots.push_back(das);
	}
	
	/* allocate space in new frame depot and map input depots*/
	Frame newFrame(to.spec());
	to.clear();
	for (size_t i = 0; i < numFrames; i++) {
		for (size_t j = 0; j < fromDepots.size(); j++) {
			const FrameDepot* curDepot = fromDepots[j].depot;
			const FrameLink* curLink = fromDepots[j].link;
			
			apply_sub_link<bool, _forward>(curDepot->bools(i), 
										   curLink->bLink, 
										   newFrame.bools);
			apply_sub_link<int, _forward>(curDepot->ints(i), 
										  curLink->iLink, 
										  newFrame.ints);
			apply_sub_link<float, _forward>(curDepot->floats(i), 
											curLink->fLink, 
											newFrame.floats);
			apply_sub_link<double, _forward>(curDepot->doubles(i), 
											 curLink->dLink, 
											 newFrame.doubles);
		}
		to.add(newFrame);
	}
}

/* perform a reverse mapping through a node's reverse frame operator */
auto_ptr<FrameDepot> create_node_input(Node& node,
									const Graph& graph,
									set<Node*>& taughtNodes,
									map<Node*, FrameDepot*>& inputDepots,
									map<Node*, FrameDepot*>& outputDepots) {
	
	/* ensure the output frame depot is prepared for the reverse mapping */
	if (outputDepots.count(&node) == 0) {
		prepare_outputs(node, graph, taughtNodes, inputDepots, outputDepots);
	}
	
	const FrameDepot* outDepot = outputDepots[&node];
	if (outDepot == NULL) {
		/* TODO: better error reporting */
		throw NullPtr();
	}
	
	auto_ptr<FrameDepot> nodeInput(new FrameDepot(node.inputSpec()));
	auto_ptr<FrameOperator> op = node.createReverseOp();
	Frame inputFrame(node.outputSpec());
	Frame outputFrame(node.inputSpec());
	
	/* map outputs through reverse operator */
	for (size_t i = 0; i < outDepot->size(); i++) {
		outDepot->fill(i, inputFrame);
		op->operate(inputFrame, outputFrame);
		nodeInput->add(outputFrame);
	}
	
	return nodeInput;
}

/* perform a forward mapping through a node's forward frame operator */
auto_ptr<FrameDepot> create_node_output(Node& node,
									 const Graph& graph,
									 set<Node*>& taughtNodes,
									 map<Node*, FrameDepot*>& inputDepots,
									 map<Node*, FrameDepot*>& outputDepots) {
	/* ensure the output frame depot is prepared for the reverse mapping */
	if (inputDepots.count(&node) == 0) {
		prepare_inputs(node, graph, taughtNodes, inputDepots, outputDepots);
	}
	
	const FrameDepot* inDepot = inputDepots[&node];
	if (inDepot == NULL) {
		/* TODO: better error reporting */
		throw NullPtr();
	}
	
	auto_ptr<FrameDepot> nodeOutput(new FrameDepot(node.inputSpec()));
	auto_ptr<FrameOperator> op = node.createForwardOp();
	Frame inputFrame(node.inputSpec());
	Frame outputFrame(node.outputSpec());
	
	/* map outputs through reverse operator */
	for (size_t i = 0; i < inDepot->size(); i++) {
		inDepot->fill(i, inputFrame);
		op->operate(inputFrame, outputFrame);
		nodeOutput->add(outputFrame);
	}
	
	return nodeOutput;
}



void prepare_inputs(Node& node,
					const Graph& graph,
					set<Node*>& taughtNodes,
					map<Node*, FrameDepot*>& inputDepots,
					map<Node*, FrameDepot*>& outputDepots) {
	/* check if input depot already exists */
	if (inputDepots.count(&node) == 1) {
		return;
	}
	
	/* force outputs of other nodes to exists */
	set<Node*> inputNodes;
	get_immediate_inputs(node, graph, inputNodes);
	for (set<Node*>::const_iterator nodeIter = inputNodes.begin();
		 nodeIter != inputNodes.end();
		 nodeIter++) {
		
		if (taughtNodes.count(*nodeIter) == 0) {
			teach_node(**nodeIter, 
					   graph, 
					   taughtNodes, 
					   inputDepots, 
					   outputDepots);
		}
		
		if (outputDepots.count(*nodeIter) == 0) {
			auto_ptr<FrameDepot> depot = create_node_output(**nodeIter,
														 graph,
														 taughtNodes,
														 inputDepots,
														 outputDepots);
			outputDepots[*nodeIter] = depot.release();
		}
	}
	
	/* create new depot */
	auto_ptr<FrameDepot> newDepot(new FrameDepot(node.inputSpec()));
	map_frame_depots<true>(graph.getInputEdges(&node), 
						   *newDepot, 
						   outputDepots);
	inputDepots[&node] = newDepot.release();
}

void prepare_outputs(Node& node,
					 const Graph& graph,
					 set<Node*>& taughtNodes,
					 map<Node*, FrameDepot*>& inputDepots,
					 map<Node*, FrameDepot*>& outputDepots) {
	/* check if input depot already exists */
	if (outputDepots.count(&node) == 1) {
		return;
	}
	
	/* force outputs of other nodes to exists */
	set<Node*> outputNodes;
	get_immediate_outputs(node, graph, outputNodes);
	for (set<Node*>::const_iterator nodeIter = outputNodes.begin();
		 nodeIter != outputNodes.end();
		 nodeIter++) {
		
		if (taughtNodes.count(*nodeIter) == 0) {
			teach_node(**nodeIter, 
					   graph, 
					   taughtNodes, 
					   inputDepots, 
					   outputDepots);
		}
		
		if (inputDepots.count(*nodeIter) == 0) {
			auto_ptr<FrameDepot> depot = create_node_input(**nodeIter,
														graph,
														taughtNodes,
														inputDepots,
														outputDepots);
			inputDepots[*nodeIter] = depot.release();
		}
	}
	
	auto_ptr<FrameDepot> newDepot(new FrameDepot(node.outputSpec()));
	map_frame_depots<false>(graph.getOutputEdges(&node), 
							*newDepot, 
							inputDepots);
	outputDepots[&node] = newDepot.release();
}

/* used to delete temporary frame depots made during the training process */
void delete_frame_depots(map<Node*, FrameDepot*>& frames) {
	for (map<Node*, FrameDepot*>::iterator iter = frames.begin();
		 iter != frames.end();
		 iter++) {
		delete iter->second;
	}
	frames.clear();
}


void get_immediate_inputs(Node& node,
						  const Graph& graph,
						  set<Node*>& inputs) {
	vector<Edge*> inputEdges = graph.getInputEdges(&node);
	for (size_t i = 0; i < inputEdges.size(); i++) {
		inputs.insert(inputEdges[i]->from);
	}
}

void get_immediate_outputs(Node& node,
						   const Graph& graph,
						   set<Node*>& inputs) {
	vector<Edge*> outputEdges = graph.getOutputEdges(&node);
	for (size_t i = 0; i < outputEdges.size(); i++) {
		inputs.insert(outputEdges[i]->to);
	}
}
